package com.flink.hbase;

import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.java.utils.ParameterTool;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;

/**
 * 广播流作业启动器
 * 演示如何使用 JDBC 数据源和 BroadcastStream 进行数据丰富化
 */
public class BroadcastJobLauncher {

    private static final Logger LOG = LoggerFactory.getLogger(BroadcastJobLauncher.class);

    public static void main(String[] args) throws Exception {
        LOG.info("启动 Flink 广播流数据丰富化作业");

        // 解析命令行参数
        ParameterTool parameterTool = ParameterTool.fromArgs(args);
        
        // 验证必需参数
        validateRequiredParameters(parameterTool);

        // 创建 Flink 执行环境
        StreamExecutionEnvironment env = createExecutionEnvironment(parameterTool);

        // 构建数据流处理管道
        buildBroadcastDataPipeline(env, parameterTool);

        // 启动作业
        String jobName = parameterTool.get("job.name", "Flink-Broadcast-Enrichment");
        LOG.info("启动作业: {}", jobName);
        env.execute(jobName);
    }

    /**
     * 验证必需参数
     */
    private static void validateRequiredParameters(ParameterTool parameterTool) {
        String[] requiredParams = {
            "jdbc.url",
            "jdbc.username", 
            "jdbc.password",
            "dimension.table.name"
        };

        for (String param : requiredParams) {
            if (!parameterTool.has(param)) {
                throw new IllegalArgumentException("缺少必需参数: " + param);
            }
        }

        LOG.info("参数验证通过");
    }

    /**
     * 创建 Flink 执行环境
     */
    private static StreamExecutionEnvironment createExecutionEnvironment(ParameterTool parameterTool) {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        // 设置全局参数
        env.getConfig().setGlobalJobParameters(parameterTool);

        // 设置并行度
        int parallelism = parameterTool.getInt("parallelism", 4);
        env.setParallelism(parallelism);
        LOG.info("设置并行度: {}", parallelism);

        // 启用检查点
        long checkpointInterval = parameterTool.getLong("checkpoint.interval", 60000);
        env.enableCheckpointing(checkpointInterval);
        LOG.info("启用检查点，间隔: {}ms", checkpointInterval);

        // 配置重启策略
        env.setRestartStrategy(
            org.apache.flink.api.common.restartstrategy.RestartStrategies.fixedDelayRestart(
                parameterTool.getInt("restart.attempts", 3),
                org.apache.flink.api.common.time.Time.milliseconds(
                    parameterTool.getLong("restart.delay", 10000)
                )
            )
        );

        return env;
    }

    /**
     * 构建广播数据处理管道
     */
    private static void buildBroadcastDataPipeline(StreamExecutionEnvironment env, 
                                                 ParameterTool parameterTool) {
        
        // 1. 创建维度数据广播流
        DataStream<DimensionData> dimensionStream = createDimensionStream(env, parameterTool);
        
        // 2. 创建主数据流（用户数据）
        DataStream<User> userStream = createUserStream(env, parameterTool);
        
        // 3. 创建广播流
        BroadcastStream<DimensionData> broadcastStream = dimensionStream
                .broadcast(UserEnrichmentFunction.DIMENSION_STATE_DESCRIPTOR);
        
        // 4. 连接主流和广播流进行数据丰富化
        SingleOutputStreamOperator<EnrichedUser> enrichedUserStream = userStream
                .connect(broadcastStream)
                .process(new UserEnrichmentFunction())
                .name("User Data Enrichment")
                .uid("user-data-enrichment");
        
        // 5. 添加水印和窗口处理（可选）
        DataStream<EnrichedUser> watermarkedStream = enrichedUserStream
                .assignTimestampsAndWatermarks(
                    WatermarkStrategy.<EnrichedUser>forBoundedOutOfOrderness(Duration.ofSeconds(5))
                        .withTimestampAssigner((event, timestamp) -> event.getTs())
                );
        
        // 6. 输出结果
        outputResults(watermarkedStream, parameterTool);
        
        LOG.info("广播流数据处理管道构建完成");
    }

    /**
     * 创建维度数据流
     */
    private static DataStream<DimensionData> createDimensionStream(StreamExecutionEnvironment env, 
                                                                 ParameterTool parameterTool) {
        
        String jdbcUrl = parameterTool.getRequired("jdbc.url");
        String username = parameterTool.getRequired("jdbc.username");
        String password = parameterTool.getRequired("jdbc.password");
        String tableName = parameterTool.getRequired("dimension.table.name");
        long intervalMs = parameterTool.getLong("dimension.refresh.interval", 300000); // 5分钟
        
        // 创建 JDBC 维度数据源
        JdbcDimensionSource dimensionSource = JdbcDimensionSource.create(
            jdbcUrl, username, password, tableName, intervalMs);
        
        return env.addSource(dimensionSource)
                  .name("JDBC Dimension Source")
                  .uid("jdbc-dimension-source")
                  .setParallelism(1); // 维度数据源通常设置为并行度1
    }

    /**
     * 创建主数据流（用户数据）
     */
    private static DataStream<User> createUserStream(StreamExecutionEnvironment env, 
                                                   ParameterTool parameterTool) {
        
        String sourceType = parameterTool.get("user.source.type", "mock");
        
        switch (sourceType.toLowerCase()) {
            case "mock":
                return createMockUserStream(env, parameterTool);
            case "hbase":
                return createHBaseUserStream(env, parameterTool);
            case "kafka":
                return createKafkaUserStream(env, parameterTool);
            default:
                throw new IllegalArgumentException("不支持的用户数据源类型: " + sourceType);
        }
    }

    /**
     * 创建模拟用户数据流
     */
    private static DataStream<User> createMockUserStream(StreamExecutionEnvironment env, 
                                                       ParameterTool parameterTool) {
        
        int mockDataRate = parameterTool.getInt("mock.data.rate", 100); // 每秒100条
        
        return env.addSource(new SourceFunction<User>() {
            private volatile boolean running = true;
            private int counter = 0;

            @Override
            public void run(SourceContext<User> ctx) throws Exception {
                while (running) {
                    User user = new User();
                    user.setUserId("user_" + String.format("%06d", counter % 10000));
                    user.setName("User " + counter);
                    user.setAge(20 + (counter % 60));
                    user.setGender(counter % 2 == 0 ? "Male" : "Female");
                    user.setTs(System.currentTimeMillis());
                    user.setDt(java.time.LocalDate.now().toString());
                    
                    ctx.collect(user);
                    counter++;
                    
                    // 控制数据生成速率
                    Thread.sleep(1000 / mockDataRate);
                }
            }

            @Override
            public void cancel() {
                running = false;
            }
        }).name("Mock User Source");
    }

    /**
     * 创建 HBase 用户数据流
     */
    private static DataStream<User> createHBaseUserStream(StreamExecutionEnvironment env, 
                                                        ParameterTool parameterTool) {
        // 这里可以复用之前创建的 HBaseAsyncSource
        String tableName = parameterTool.get("hbase.table.name", "user");
        int asyncPoolSize = parameterTool.getInt("async.pool.size", 10);
        long timeoutMs = parameterTool.getLong("async.timeout.ms", 5000);
        
        // 创建用户ID流
        DataStream<String> userIdStream = env.addSource(new SourceFunction<String>() {
            private volatile boolean running = true;
            private int counter = 0;

            @Override
            public void run(SourceContext<String> ctx) throws Exception {
                while (running) {
                    ctx.collect("user_" + String.format("%06d", counter % 10000));
                    counter++;
                    Thread.sleep(100); // 每秒10个用户ID
                }
            }

            @Override
            public void cancel() {
                running = false;
            }
        });

        // 使用异步 HBase 查询
        HBaseAsyncSource hbaseSource = new HBaseAsyncSource(tableName, asyncPoolSize, timeoutMs);
        
        return org.apache.flink.streaming.api.datastream.AsyncDataStream.unorderedWait(
            userIdStream,
            hbaseSource,
            timeoutMs,
            java.util.concurrent.TimeUnit.MILLISECONDS,
            asyncPoolSize
        ).name("HBase User Source");
    }

    /**
     * 创建 Kafka 用户数据流
     */
    private static DataStream<User> createKafkaUserStream(StreamExecutionEnvironment env, 
                                                        ParameterTool parameterTool) {
        // TODO: 实现 Kafka 用户数据源
        throw new UnsupportedOperationException("Kafka 用户数据源尚未实现");
    }

    /**
     * 输出处理结果
     */
    private static void outputResults(DataStream<EnrichedUser> enrichedUserStream, 
                                    ParameterTool parameterTool) {
        
        String outputType = parameterTool.get("output.type", "print");
        
        switch (outputType.toLowerCase()) {
            case "print":
                enrichedUserStream
                    .print("Enriched User")
                    .name("Print Enriched Users");
                break;
                
            case "file":
                String outputPath = parameterTool.get("output.path", "/tmp/enriched-users");
                enrichedUserStream
                    .writeAsText(outputPath)
                    .name("File Output");
                break;
                
            case "kafka":
                // TODO: 实现 Kafka 输出
                throw new UnsupportedOperationException("Kafka 输出尚未实现");
                
            case "elasticsearch":
                // TODO: 实现 Elasticsearch 输出
                throw new UnsupportedOperationException("Elasticsearch 输出尚未实现");
                
            default:
                LOG.warn("未知的输出类型: {}, 使用默认的 print 输出", outputType);
                enrichedUserStream.print("Enriched User");
        }
    }

    /**
     * 打印使用说明
     */
    public static void printUsage() {
        System.out.println("Flink 广播流数据丰富化作业使用说明:");
        System.out.println();
        System.out.println("必需参数:");
        System.out.println("  --jdbc.url <url>                    JDBC 数据库连接 URL");
        System.out.println("  --jdbc.username <username>          数据库用户名");
        System.out.println("  --jdbc.password <password>          数据库密码");
        System.out.println("  --dimension.table.name <table>      维度表名称");
        System.out.println();
        System.out.println("可选参数:");
        System.out.println("  --job.name <name>                   作业名称");
        System.out.println("  --parallelism <num>                 并行度");
        System.out.println("  --user.source.type <type>           用户数据源类型 (mock|hbase|kafka)");
        System.out.println("  --dimension.refresh.interval <ms>   维度数据刷新间隔(毫秒)");
        System.out.println("  --mock.data.rate <rate>             模拟数据生成速率(每秒记录数)");
        System.out.println("  --output.type <type>                输出类型 (print|file|kafka)");
        System.out.println("  --output.path <path>                输出路径(文件输出时)");
        System.out.println();
        System.out.println("示例:");
        System.out.println("  java -cp flink-hbase-etl.jar com.flink.hbase.BroadcastJobLauncher \\");
        System.out.println("    --jdbc.url jdbc:mysql://localhost:3306/test \\");
        System.out.println("    --jdbc.username root \\");
        System.out.println("    --jdbc.password password \\");
        System.out.println("    --dimension.table.name user_dimensions \\");
        System.out.println("    --user.source.type mock \\");
        System.out.println("    --output.type print");
    }
} 